# Filter Events
import asyncio
from langchain.chat_models import init_chat_model
from langchain_core.output_parsers import JsonOutputParser
import os

key = os.getenv("OPENAI_API_KEY")
# print(key)
api_key = str(key)

model = init_chat_model(
    model="gpt-4o-mini",
    model_provider="openai",
    base_url="https://api.zetatechs.com/v1",
    api_key=api_key
)
#  这是通过name来进行分类
#  You can filter by either component name, component tags or component type.
chain = model.with_config({"run_name": "model"}) | JsonOutputParser().with_config(
    {"run_name": "my_parser"}
)


async def use():
    max_events = 0
    async for event in chain.astream_events(
            'output a list of the countries france, spain and japan and their populations in JSON format. Use a dict with an outer key of "countries" which contains a list of countries. Each country should have the key `name` and `population`',
            include_types=["chat_model"],
    ):
        print(event)
        max_events += 1
        if max_events > 10:
            # Truncate output
            print("...")
            break


# asyncio.run(use())

chain = (model | JsonOutputParser()).with_config({"tags": ["my_chain"]})


async def use1():
    max_events = 0
    async for event in chain.astream_events(
            'output a list of the countries france, spain and japan and their populations in JSON format. Use a dict with an outer key of "countries" which contains a list of countries. Each country should have the key `name` and `population`',
            include_tags=["my_chain"],
    ):
        print(event)
        max_events += 1
        if max_events > 10:
            # Truncate output
            print("...")
            break


asyncio.run(use1())
